from typing import List, Tuple
import os
from langchain_community.callbacks.fiddler_callback import RUN_ID
from research_agent.core.propose_question import QuestionProposer
from research_agent.core.resolve_HYDE import HYDEResolver
from research_agent.core.paper_review import PaperReviewer
from research_agent.core.propose_suggestion import SuggestionProposer
from research_agent.core.summary import SummaryExtractor
from research_agent.core.judge_topic_type import TopicJudger
from research_agent.core.writer import Writer
from research_agent.core.pipeline_reference import CitationPipeline
from research_agent.core.write_outline import OutlineWriter
from research_agent.core.generate_hyde import HYDEGenerator
from datetime import datetime
import logging
from logging import Logger
from pyaml_env import parse_config
from pathlib import Path
import asyncio
import pickle
from research_agent.core.run_xgboost import Run_xgb
from functools import wraps
from research_agent.core.config import Config
from research_agent.core.generate_embedding import EmbeddingGenerator
from jinja2 import Environment

logger = logging.getLogger(__name__)
def async_retry(retries=3, delay=1):
    """
    异步重试装饰器
    Args:
        retries (int): 最大重试次数
        delay (int): 重试间隔时间(秒)
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(retries):
                try:
                    result = await func(*args, **kwargs)
                    return result
                except Exception as e:
                    last_exception = e
                    if attempt < retries - 1:
                        await asyncio.sleep(delay)
                    logging.warning(f"第 {attempt + 1} 次尝试失败: {str(e)}")
            raise last_exception
        return wrapper
    return decorator



class Pipeline:
    """A pipeline for generating and iteratively improving research survey papers.
    
    This class orchestrates the process of generating research surveys by:
    1. Proposing research questions
    2. Resolving these questions with relevant context
    3. Writing draft sections
    4. Reviewing and improving the paper through iterations
    
    Attributes:
        iteration_limit (int): Maximum number of improvement iterations
        score_threshold (float): Minimum acceptable quality score
    """
    def __init__(self, draft_iteration_output_dir:Path=None,xgb_model_output_dir:Path=None):
        """Initialize the pipeline components.
        
        Args:
            logger (Logger, optional): Custom logger instance. If None, creates a default logger.
            output_dir (str, optional): Output directory path for generated files. Defaults to "output".
        """

        self.propose_question = QuestionProposer()
        self.resolve_HYDE = HYDEResolver(top_k = 5)
        self.propose_suggestion = SuggestionProposer()
        self.writer = Writer()
        self.paper_reviewer = PaperReviewer()
        self.HYDE_generator = HYDEGenerator()
        self.outline_writer = OutlineWriter()
        self.topic_judger = TopicJudger()
        self.chunk_summarizer = SummaryExtractor()
        self.citation_processor = None
        self.xgb_model = None
        self.iteration_limit = 1
        self.score_threshold = 9.5
        self.context = []
        self.related_papers = []
        self.context_embedding = None
        absolute_path = os.path.abspath(Config.YAML_CONFIG)
        configs = parse_config(absolute_path)
        self.embedding_gen = EmbeddingGenerator()
        # Setup logging
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        draft_iteration_output_dir = Path(draft_iteration_output_dir / f"{timestamp}")
        draft_iteration_output_dir.mkdir(exist_ok=True)
        self.draft_iteration_output_dir = draft_iteration_output_dir
        self.context_summary = None

        # 这里得加载一个xgboost,用于判断topic属于四种综述中的哪种
        #with open(xgb_model_output_dir/"xgb_model.data","rb") as file:
        #    self.xgb_model = pickle.load(file)

    async def _process_results(self, results: List[dict]) -> Tuple[List[str], List[str]]:
        """Process question resolution results into context and related papers.
        
        Args:
            results: List of resolution results containing context and related papers
            
        Returns:
            Tuple containing processed context string and related papers string
        """
        try:
            context = []
            related_papers = []
            
            for result in results:
                if not isinstance(result, dict) or 'context' not in result or 'related_papers' not in result:
                    logger.warning(f"Malformed result detected: {result}")
                    continue
                context.extend(result["context"])
                related_papers.extend(result["related_papers"])
            return context, related_papers
        except Exception as e:
            logger.error(f"Error processing results: {str(e)}")
            raise

    async def write_a_draft_survey(self, topic: str,survey_type:int,background_questions) -> str:
        """Generate initial draft sections for a research survey.
        
        Args:
            topic: The research topic to write about
            survey_type: The type of survey to write
            
        Returns:
            Tuple of (introduction, related_work, conclusion) sections
            
        Raises:
            ValueError: If topic is empty or invalid
        """
        if not topic.strip():
            raise ValueError("Topic cannot be empty")
            
        logger.info(f"Starting draft survey generation for topic: {topic}\n ")
        try:
            # 使用重试装饰器包装关键操作
            @async_retry(retries=3, delay=1)
            async def generate_questions(survey_type:int):
                return await self.propose_question.propose_question(topic,survey_type)

            @async_retry(retries=3, delay=1)
            async def resolve_HYDEs(hydes):
                return await self.resolve_HYDE.resolve_HYDEs(hydes)

            @async_retry(retries=3, delay=1)
            async def generate_HYDEs(questions):
                return await self.HYDE_generator.generate_HYDEs(questions)

            @async_retry(retries=3, delay=1)
            async def generate_outline(survey_type,context, related_papers):
                return await self.outline_writer.write_outline(topic, survey_type,context, related_papers)

            @async_retry(retries=3, delay=1)
            async def generate_initial_draft(outline, context,context_embedding,related_papers):
                return await self.writer.write_initial_draft(topic, outline, context, context_embedding,related_papers)

            # 执行带有重试机制的操作
            questions = await generate_questions(survey_type)

            logger.info(f"Generated {len(questions)} questions")
            questions.extend(background_questions)
            logger.info(f"There are {len(questions)} questions in total")
            hydes = await generate_HYDEs(questions)
            logger.info(f"Generated {len(hydes)} HYDEs")
            results = await resolve_HYDEs(hydes)
            logger.info(f"{len(results)} HYDEs are resolved")
            #self.logger.debug(f"Generated results: {results}")
            self.context, self.related_papers = await self._process_results(results)
            # 获取context的summary
            self.context_summary = await self.chunk_summarizer.summarize_chunks(self.context,topic,self.writer.get_language())
            logger.info(f"Generated {len(self.context_summary)} context summary")
            # 我需要把context放入向量数据库中
            # 获取每个paper的summary的embedding,并记录失败的paper在队列中的序号
            context_embedding_tuple,failed_id = await self.embedding_gen.generate_embeddings(self.context)

            # 将成功的paper留在self.context和self.related_papers中
            # 将failed_id是一个List,将其中的元素前后对调一下
            logger.info(f"Generated {len(context_embedding_tuple)} context embeddings,with {len(failed_id)} context went wrong,the failed id is {failed_id}")
            self.context = [context_i for _,context_i in context_embedding_tuple]
            self.context_embedding = [context_embedding_i for context_embedding_i,_ in context_embedding_tuple]
            for i in failed_id[::-1]:
                del self.related_papers[i]
                del self.context_summary[i]

            outline = await generate_outline(survey_type,'\n\n'.join(self.context_summary), '\t'.join(self.related_papers))
            logger.info(f"Successfully generated initial outline:{outline}")

            initial_draft = await generate_initial_draft(outline, self.context, self.context_embedding,self.related_papers)

            logger.info("Successfully generated initial draft")
            return initial_draft
            
        except Exception as e:
            logger.error(f"Error in draft generation: {str(e)}")
            raise

    async def paper_and_rewrite(self, topic: str, raw_draft: str,review_questions: list[str],survey_type:int):
        try:
            #context, related_papers = [], set()
            @async_retry(retries=3, delay=1)
            async def resolve_HYDEs(hydes):
                return await self.resolve_HYDE.resolve_HYDEs(hydes)

            @async_retry(retries=3, delay=1)
            async def generate_HYDEs(review_questions):
                return await self.HYDE_generator.generate_HYDEs(review_questions)

            @async_retry(retries=3, delay=1)
            async def propose_suggestions():
                return await self.propose_suggestion.propose_suggestion(
                    topic, raw_draft, context, related_papers
                )

            @async_retry(retries=3, delay=1)
            async def rewrite_draft(suggestions):
                return await self.writer.rewrite_draft(
                    topic, self.context, self.context_embedding,suggestions, self.related_papers, raw_draft
                )



            hydes = await generate_HYDEs(review_questions)
            results = await resolve_HYDEs(hydes)
            for result in results:
                self.context.append(result["context"])
                self.related_papers.append(result["related_papers"])

            context = "\n".join(self.context)
            related_papers = "\t".join(self.related_papers)

            suggestions = await propose_suggestions()
            logger.debug(f"Generated suggestions: {suggestions}")
            suggestions = "\n".join(suggestions)

            updated_draft = await rewrite_draft(suggestions)
            return updated_draft

        except Exception as e:
            logger.error(f"Error in paper rewrite: {str(e)}")
            raise

    async def review_draft(self, topic: str, raw_draft: str):
        try:
            # context, related_papers = [], set()
            @async_retry(retries=3, delay=1)
            async def review_paper(topic, paper_draft):
                return await self.paper_reviewer.review_paper(topic, paper_draft, self.writer.survey)

            survey = await review_paper(topic, raw_draft)
            return survey
        except Exception as e:
            logger.error(f"Error in paper rewrite: {str(e)}")
            raise

    async def iteration(self, topic: str,language:str) :
        """Iteratively improve the survey paper until quality threshold is met.
        
        Args:
            topic: The research topic
            
        Returns:
            Tuple of final (introduction, related_work, conclusion) sections
        """
        if not topic.strip():
            raise ValueError("Topic cannot be empty")
        # 先用xgboost判断topic属于哪种综述
        #survey_type = self.xgb_model.predict(topic)

        self.writer.set_language(language)
        self.paper_reviewer.set_language(language)
        self.outline_writer.set_language(language)

        self.topic_judger.set_language(language)

        survey_type_info = await self.topic_judger.judge_topic_type(topic)
        survey_type = int(survey_type_info["type_judgment"]["type"])
        background_questions = survey_type_info["background_questions"]
        logger.info(f"There are 4 type of survey, namely: \n1.对一个技术概念的调研综述；\n"
                         f"2.对一个方向研究现状的综述；\n3.对多个方法的对比分析综述；\n4.对一个技术方法研究脉络的综述。\n"
                         f"The topic of \"{topic}\" belongs to {survey_type}")
        logger.info(f"Generated {len(background_questions)} background questions")
        logger.info(f"Starting iterative improvement for topic: {topic}")
        '''
        {
            "revised_topic": "string",
            "type_judgment": {
                "type": "number (1-4)",
                "reason": "string"
            },
            "background_questions": ["string"]
        }
        '''
        try:
            paper_draft = await self.write_a_draft_survey(topic,survey_type,background_questions)
            score = 0
            iteration = 0
            # 写入Markdown文件
            with open(self.draft_iteration_output_dir/f'{iteration}.md', 'w', encoding='utf-8') as file:
                file.write(paper_draft)
            #while score < self.score_threshold and iteration < self.iteration_limit:
            while iteration < self.iteration_limit:
                logger.info(f"Starting iteration {iteration + 1}/{self.iteration_limit}")
                # review的时候直接就把格式的问题修改了，包括引用句锚点、表格、插图、重复内容等的删除
                # 然后返回每个小节需要补充的问题
                self.writer.survey = await self.review_draft(topic, paper_draft)
                paper_draft = self.writer.survey.get_paper_content()
                #score = review["score"]
                #self.logger.info(f"Review score: {score}/{self.score_threshold}")
                #self.logger.debug(f"Review questions: {review['questions']}")
                #if score < self.score_threshold:
                    # 在rewrite时候，针对每个小节提出的问题，进行HYDE和RAG,然后重写小节内容
                #    paper_draft = await self.paper_and_rewrite(
                #        topic, paper_draft,review["questions"],survey_type
                #    )
                iteration += 1
                # 写入Markdown文件
                with open(self.draft_iteration_output_dir / f'{iteration}.md', 'w', encoding='utf-8') as file:
                    file.write(paper_draft)
            logger.info(f"Completed iterations with final score: {score}")
            # 添加引言和总结内容
            # 先从self.writer.survey中获取非引言和总结的内容

            max_section_code = self.writer.survey.get_max_section_code()
            outline = self.writer.survey.get_whole_outline()
            paper_draft = self.writer.survey.get_paper_content(except_sections=[str(1), max_section_code])
            # 基于这些内容生成引言和总结内容
            tasks = [self.writer.write_title(topic, outline,paper_draft),self.writer.write_introduction(topic, outline,paper_draft),self.writer.write_conclusion(topic,outline,paper_draft)]

            paper_title,introduction,conclusion = await asyncio.gather(*tasks, return_exceptions=True)
            # 生成文章全部内容
            paper_draft = f"## 1 Introduction\n\n{introduction}\n\n {paper_draft} \n\n## {max_section_code} Conclusion\n\n{conclusion}"
            # 为文章内容添加参考文献
            self.citation_processor = CitationPipeline(topic =topic,content = paper_draft,draft_iteration_output_dir=self.draft_iteration_output_dir)
            paper_draft = await self.citation_processor.pipeline_reference()
            paper_draft = f"# {paper_title}\n\n{paper_draft}"
            return paper_draft
            
        except Exception as e:
            logger.error(f"Error in iteration process: {str(e)}")
            raise

